Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 728-745: The chunked destination slicing in transfer logic is now
using chunk offsets, but the token alignment still assumes each chunk maps to a
suffix ending at token_range.end. Update the code around the chunked path in
transfer.py and the downstream token-start calculation to derive starts from
chunk_block_offset, or require callers to provide per-chunk KVSlice.token_range
for each chunk. Make sure the block selection and token-range alignment stay
consistent for prefix-cache and SWA cases so the written blocks match the
intended chunk.
- Around line 523-524: The abort/result notification path in transfer.py still
uses write_meta.slice_id, which can conflict with the receiver’s single-task
slice handling. Update the abort send logic in the relevant transfer routine to
mirror the success path by reporting receiver_slice_id as 0 for aborts too, so
the receiver does not see a later-chunk slice ID and hit its slice assertion.
Keep the existing task/event unblocking behavior intact while ensuring the
aborted/failure result is always sent to receiver slice 0.
- Around line 736-743: The chunk-to-destination mapping in transfer.py is too
strict for exhausted layer groups: when len(src_block_ids) is 0, the current
bounds check in the chunk slicing logic still raises on advanced chunk_offset
values. Update the chunk handling around the dst_block_ids slice so empty source
chunks become a no-op and do not trigger the out-of-bounds error; keep the
existing bounds validation for non-empty chunks in the same chunk
offset/full_dst_block_ids path.

In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Line 648: The `respond_and_send_async` skip guard in `transceiver.py` needs
both a lint fix and a logic check: move the `return` onto its own line to
satisfy E701, and verify the condition around `rid in self._send_sessions and
rid not in self._pipelined_chunk_offsets` correctly prevents duplicate sends
while pipelined chunks are still outstanding. If needed, adjust the guard so the
full `_create_kv_slices` resend path only runs when it is safe, using the
existing `_send_sessions` and `_pipelined_chunk_offsets` state to avoid
duplicate transfer.
- Around line 591-611: The chunking logic in send_prefill_chunk() and the
_pipelined_chunk_offsets update can split KV slices on token boundaries that are
not aligned to tokens_per_block, which causes the boundary block to be resent
and offsets to drift. Adjust the prefill chunk selection so every chunk boundary
lands on a KV block boundary (or clamp the sliding-window fallback so it only
overlaps when it evenly divides tokens_per_block), and then recompute
_pipelined_chunk_offsets from the actual block count in the chunk.

In `@tests/unittest/disaggregated/test_kv_transfer.py`:
- Around line 1789-1825: The send/receive flow is using the wrong API shape:
TxSession.send() and RxSession.receive() should be called with a fully populated
KVSlice rather than extra kwargs, and they do not return futures. Update the
test setup around KVSlice, sender_session.send(), and
receiver_sessions/RxSession.receive() to set chunk_block_offset and cuda_event
on the slice object before calling send/receive, then replace the .result()
waits with wait_complete()/wait_complete(blocking=True) on the session or slice
as appropriate.

---

Outside diff comments:
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 488-507: The task transition in transfer.py is happening too late
in the prefill-transfer flow: `task._slice.cuda_event.synchronize()` runs while
the task is still `INIT`, so `cancel_request()` can miss it and free KV pages
too early. In the transfer path around `task`, `session.lock`, and
`TaskStatus.TRANSFERRING`, move the INIT→TRANSFERRING state update (with the
session ERROR/CANCELLED abort check) before waiting on the CUDA event, and keep
the abort branch ahead of synchronization so in-flight work is visible before
any blocking wait.

---

Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Around line 582-585: The send-session initialization in transceiver logic has
a redundant assignment because _get_or_create_send_session already stores the
session in self._send_sessions. Update the rid-not-in-self._send_sessions branch
in transceiver.py to follow the same pattern as respond_and_send_async by simply
invoking _get_or_create_send_session(req) for its side effects, then keep
setting _ever_had_send_session and _pipelined_chunk_offsets[rid] as before.
- Around line 602-604: The pipelined-send path in transceiver.py still contains
unresolved correctness TODOs in send_prefill_chunk and respond_and_send_async,
especially around token_range, mamba_state_index, layer_range, req.state
transitions, and the offset accumulation logic. Resolve these TODO(athenac)
questions by verifying the metadata semantics, fixing the offset calculation,
and making the state update explicit and correct before merge. Also remove the
duplicated logic between send_prefill_chunk and respond_and_send_async by
consolidating the shared send/metadata assembly into a single helper so the two
paths stay consistent.